'''
Author: SUNNY
Description: 
Date: 2025-11-04 18:29:26
'''
import os
from openai import OpenAI

class DeepSeekStreamClient:
    def __init__(self):
        self.client = OpenAI(
            api_key=os.environ.get("AI_DEEPSEEK_API_KEY"),
            base_url="https://api.deepseek.com"
        )
        self.messages = [
            {"role": "system", "content": "You are a helpful assistant"}
        ]

    def chat_stream(self, user_input):
        # 添加用户消息
        self.messages.append({"role": "user", "content": user_input})

        try:
            # 创建流式响应
            stream = self.client.chat.completions.create(
                model="deepseek-chat",
                messages=self.messages,
                stream=True,
                max_tokens=2048,
                temperature=0.7
            )

            print("DeepSeek 思考中...\n", end="", flush=True)
            assistant_response = ""

            # 处理流式输出
            for chunk in stream:
                if (chunk.choices[0].delta.content is not None and chunk.choices[0].delta.content != ""):
                    content = chunk.choices[0].delta.content
                    print(content, end="", flush=True)
                    assistant_response += content

            print()  # 换行

            # 将助手回复添加到消息历史
            self.messages.append({"role": "assistant", "content": assistant_response})

            return assistant_response
        
        except Exception as e:
            print(f"\n请求失败: {e}")
            # 移除失败的用户消息
            self.messages.pop()
            return None
        
    def clear_history(self):
        """清空对话历史（保留系统消息）"""
        self.messages = [self.messages[0]]
        print("对话历史已清除")

# 交互式聊天
if __name__ == "__main__":
    import sys

    # 检查API密钥
    if not os.environ.get("AI_DEEPSEEK_API_KEY"):
        print("错误: 请设置 DEEPSEEK_API_KEY 环境变量")
        sys.exit(1)

    client = DeepSeekStreamClient()

    print("DeepSeek 流式聊天助手")
    print("输入 'clear' 清空对话历史，输入 'bye' 退出程序")
    print("=" * 50)


    while True:
        try:
            user_input = input("\n用户: ").strip()

            if not user_input:
                continue

            if user_input.lower() == "bye":
                print("再见！")
                break
            elif user_input.lower() == "clear":
                client.clear_history()
                continue

            # 流式聊天
            client.chat_stream(user_input)

        except KeyboardInterrupt:
            print("\n\n程序被用户中断")
            break
        
        except Exception as e:
            print(f"\n发生错误: {e}")